import os
import shutil
from sklearn.model_selection import train_test_split


def split_image_dataset(image_dir, output_dir, train_size=0.7, val_size=0.15, test_size=0.15, random_state=None):
    """
    将图片数据集划分为训练集、验证集和测试集。

    :param image_dir: str, 原始图片文件夹路径
    :param output_dir: str, 输出文件夹路径
    :param train_size: float, 训练集所占比例
    :param val_size: float, 验证集所占比例
    :param test_size: float, 测试集所占比例
    :param random_state: int, 随机种子
    """
    # 检查比例是否正确
    assert train_size + val_size + test_size == 1, "训练集、验证集、测试集的概率总和为1"

    # 获取所有图片文件的列表
    all_images = [os.path.join(image_dir, img) for img in os.listdir(image_dir) if
                  os.path.isfile(os.path.join(image_dir, img))]

    # 首先将图片划分为训练集+验证集 和 测试集
    train_val_images, test_images = train_test_split(all_images, test_size=test_size, random_state=random_state)

    # 计算训练集和验证集的比例
    val_relative_size = val_size / (train_size + val_size)

    # 将训练集+验证集再划分为训练集和验证集
    train_images, val_images = train_test_split(train_val_images, test_size=val_relative_size,
                                                random_state=random_state)

    # 创建输出文件夹
    os.makedirs(output_dir, exist_ok=True)
    train_dir = os.path.join(output_dir, 'train')
    val_dir = os.path.join(output_dir, 'val')
    test_dir = os.path.join(output_dir, 'test')

    os.makedirs(train_dir, exist_ok=True)
    os.makedirs(val_dir, exist_ok=True)
    os.makedirs(test_dir, exist_ok=True)

    # 复制图片到对应的文件夹
    for img in train_images:
        shutil.copy(img, train_dir)
    for img in val_images:
        shutil.copy(img, val_dir)
    for img in test_images:
        shutil.copy(img, test_dir)

    print("图片数据集划分完成：")
    print(f"训练集图片数: {len(train_images)}")
    print(f"验证集图片数: {len(val_images)}")
    print(f"测试集图片数: {len(test_images)}")


# 示例用法
if __name__ == "__main__":
    # 设置原始图片文件夹路径和输出文件夹路径
    image_dir = './image_dataset'
    output_dir = './output_directory'

    # 划分图片数据集
    split_image_dataset(image_dir, output_dir, train_size=0.7, val_size=0.15, test_size=0.15, random_state=42)
